import uuid
from typing import TypedDict, Annotated, Literal, Dict, Callable, Any, Union, Optional, Sequence
from datetime import datetime
import asyncio
import os

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START, END
from langgraph.graph import StateGraph
from langgraph.types import interrupt, Command
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from rich.console import Console
from rich.prompt import Prompt
from rich.panel import Panel
from rich.markdown import Markdown
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 初始化 Rich Console
console = Console()

# 定义中断数据类型
class InterruptData(TypedDict):
    # 中断类型
    type: str
    
    # 中断描述
    description: str

    # 中断时携带的数据
    data: Dict[str, Any]

# 创建中断数据的辅助函数
def create_interrupt(interrupt_type: str, description: str, data: Dict[str, Any] = None) -> InterruptData:
    """
    创建标准化的中断数据
    
    Args:
        interrupt_type: 中断类型
        description: 中断描述，在handle_interrupt中显示
        data: 额外的数据
    
    Returns:
        InterruptData: 格式化的中断数据
    """
    return {
        "type": interrupt_type,
        "description": description,
        "data": data or {}
    }
# 定义状态
class State(TypedDict):
    # 文本内容
    text: str
    # 对话历史
    messages: Annotated[Sequence[BaseMessage], add_messages]
    # 修订历史
    revision_history: list[dict]
    # 当前步骤
    current_step: str
    # 是否完成
    is_complete: bool
    # 润色提示
    polish_hint: str | None

# 中断处理器注册表
interrupt_handlers = {}

# 注册中断处理器函数
def register_interrupt_handler(interrupt_type: str, handler: Callable[[InterruptData], dict]):
    """
    注册一个中断处理器
    
    Args:
        interrupt_type: 中断类型
        handler: 处理函数，接收interrupt_data，返回command_data
    """
    interrupt_handlers[interrupt_type] = handler

# 默认中断处理器
def default_interrupt_handler(interrupt_data: InterruptData) -> dict:
    """
    默认的中断处理函数，用于处理未注册的中断类型
    
    Args:
        interrupt_data: 中断数据
    
    Returns:
        command_data: 处理后的指令数据
    """
    console.print(f"[yellow]未知中断类型: {interrupt_data.get('type', '未知')}[/yellow]")
    console.print("\n[bold yellow]请输入响应：[/bold yellow]")
    user_input = Prompt.ask("响应")
    return {'user_input': user_input}

# AI 节点：分析文本并提供建议
def ai_analysis_node(state: State):
    # 创建LLM
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
    
    # 获取润色提示（如果有的话）
    polish_hint = ""
    if 'polish_hint' in state and state['polish_hint']:
        polish_hint = f"\n润色提示: {state['polish_hint']}"
    
    analysis_prompt = f"""
    请润色以下文本，使其更加流畅、专业和有表现力：
    
    文本：{state['text']}
    -----------
    {polish_hint}
    -----------
    
    请直接给出润色后的完整文本。不用其他说明，默认不超过50字。
    """
    
    # 准备消息
    messages = state.get('messages', [])
    
    # 如果消息为空，创建初始消息
    if not messages:
        messages = [HumanMessage(content=analysis_prompt)]
    else:
        # 添加当前请求到消息历史
        messages.append(HumanMessage(content=analysis_prompt))
    
    # 调用LLM
    response = llm.invoke(messages)
    
    # 将LLM响应添加到消息历史
    messages.append(response)
    
    # 初始化新状态
    new_state = {"messages": messages}

    # 清除上一轮的润色提示（如果有）
    if 'polish_hint' in state:
        new_state['polish_hint'] = None

    return new_state

# 人工审核节点
def human_review_node(state: State):
    # 使用辅助函数创建中断数据
    interrupt_data = create_interrupt(
        interrupt_type="review",
        description="进入人工审核环节",
        data={
            "text": state['text'],
            "ai_suggestions": state['messages'][-1].content if state.get('messages') else None
        }
    )
    
    # 中断以获取人工输入，对中断的处理已经集中在handle_interrupt函数中
    return interrupt(interrupt_data)

# 最终确认节点
def finalize_node(state: State):
    # 使用辅助函数创建中断数据
    interrupt_data = create_interrupt(
        interrupt_type="final_confirm",
        description="进入最终确认环节",
        data={
            "text": state['text']
        }
    )
    
    # 最终确认，对中断的处理已经集中在handle_interrupt函数中
    return interrupt(interrupt_data)

# 路由函数
def route_next_step(state: State) -> Union[Literal["ai_analysis", "human_review", "finalize"], str]:
    if state.get('is_complete', False):
        return END
    return state.get('current_step', 'ai_analysis')
    
# 构建图
def build_hitl_graph():
    graph_builder = StateGraph(State)
    
    # 添加节点
    graph_builder.add_node("ai_analysis", ai_analysis_node)
    graph_builder.add_node("human_review", human_review_node)
    graph_builder.add_node("finalize", finalize_node)
    
    # 添加边
    graph_builder.add_edge(START, "ai_analysis")
    
    # AI分析后直接进入人工审核
    graph_builder.add_edge("ai_analysis", "human_review")
    
    # 人工审核节点的条件路由
    graph_builder.add_conditional_edges(
        "human_review",
        route_next_step,
        {
            "ai_analysis": "ai_analysis",
            "finalize": "finalize",
            END: END
        }
    )
    
    # 最终确认节点的条件路由
    graph_builder.add_conditional_edges(
        "finalize",
        route_next_step,
        {
            "human_review": "human_review", 
            "ai_analysis": "ai_analysis",  # 添加可能的返回到AI分析的路径
            END: END
        }
    )
    
    # 编译图
    checkpointer = InMemorySaver()
    return graph_builder.compile(checkpointer=checkpointer)

# 通用中断处理函数
async def handle_interrupt(graph, interrupt_result, config):
    """
    处理工作流中的中断
    
    Args:
        graph: LangGraph实例
        interrupt_result: 包含中断信息的结果
        config: 图配置
    
    Returns:
        更新后的状态
    """
    if '__interrupt__' not in interrupt_result:
        return interrupt_result
    
    # 获取中断数据并转换为InterruptData类型
    interrupt_data: InterruptData = interrupt_result['__interrupt__'][0].value
    interrupt_type = interrupt_data.get('type', '')
    
    # 显示中断描述信息
    description = interrupt_data.get('description', f'进入{interrupt_type}中断处理流程')
    console.print(f"\n*****[bold cyan]{description}[/bold cyan]*****")
    
    # 查找对应的处理器
    handler = interrupt_handlers.get(interrupt_type, default_interrupt_handler)
    
    # 调用处理器处理中断
    command_data = handler(interrupt_data)
    
    # 恢复执行
    response = await graph.ainvoke(
        Command(resume=command_data),
        config
    )
    
    # 检查是否还有后续中断
    if '__interrupt__' in response:
        # 递归处理后续中断
        return await handle_interrupt(graph, response, config)

    return response

# 交互式运行函数
async def run_hitl_workflow():
    graph = build_hitl_graph()
    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
    
    # 获取初始文本
    initial_text = Prompt.ask("[bold cyan]请输入要润色的文本[/bold cyan]")
    #polish_hint = Prompt.ask("[bold cyan]请输入润色提示[/bold cyan]")
    polish_hint = '美化文本，增加表现力和流畅度'
    
    # 初始状态
    state = {
        "text": initial_text,
        "messages": [],
        "revision_history": [],
        "current_step": "ai_analysis",
        "is_complete": False,
        "polish_hint": polish_hint
    }
    
    try:
        result = await graph.ainvoke(state, config)
            
        # 使用通用中断处理函数处理所有类型的中断
        state = await handle_interrupt(graph, result, config)
            
    except KeyboardInterrupt:
        console.print("\n[red]流程被用户中断[/red]")
    
    console.print("\n[bold green]✨ 流程完成！[/bold green]")
    console.print(f"[bold]最终文本：[/bold] {state.get('text', 'N/A')}")

# 注册review中断处理器
def register_review_interrupt_handler():
    def review_handler(interrupt_data: InterruptData) -> dict:
        # 获取中断的原始文本内容
        original_text = interrupt_data.get('data', {}).get('text', '')
        ai_suggestions = interrupt_data.get('data', {}).get('ai_suggestions', '')

        # 打印原始文本和AI建议
        # console.print(Panel(f"[bold]原始文本:[/bold]\n{original_text}", style="blue"))
        console.print(Panel(f"[bold]AI润色后的文本:[/bold]\n{ai_suggestions}", style="green")) 
        
        # 在处理函数中显示提示信息
        console.print("\n[bold yellow]请选择操作：\n1. 接受当前文本\n2. 给予提示后重新润色\n3. 结束流程[/bold yellow]")
        
        choice = Prompt.ask("选择", choices=['1', '2', '3'], default='1')
        command_data = {}
        
        if choice == '2':
            polish_hint = Prompt.ask("[bold cyan]请输入润色提示(可选，直接回车跳过)[/bold cyan]", default="")
            if polish_hint:
                command_data.update({
                    'reason': "根据提示润色",
                    'polish_hint': polish_hint
                })
            else:
                command_data.update({
                    'reason': "重新润色" 
                })
            
            # 重新润色时，直接设置下一步为AI分析
            command_data['current_step'] = "ai_analysis"
        elif choice == '1':
            # 接受当前文本，进入最终确认阶段
            command_data['current_step'] = "finalize"
            # 如果用户接受当前文本，将AI的建议作为新的文本
            command_data['text'] = ai_suggestions
        elif choice == '3':
            # 结束流程
            command_data['is_complete'] = True
            command_data['current_step'] = "end"
        
        return command_data
    
    register_interrupt_handler('review', review_handler)

# 注册final_confirm中断处理器
def register_final_confirm_interrupt_handler():
    def final_confirm_handler(interrupt_data: InterruptData) -> dict:
        # 获取中断的文本内容和修订历史
        text = interrupt_data.get('data', {}).get('text', '')
        revision_history = interrupt_data.get('data', {}).get('revision_history', [])
        
        # 显示最终文本
        console.print(Panel(f"[bold]AI建议的文本:[/bold]\n{text}", style="green"))  
        
        # 显示修订历史
        if revision_history:
            console.print(Panel("[bold]修订历史:[/bold]", style="blue"))
            for revision in revision_history:
                console.print(f"- {revision.get('timestamp', '')}: {revision.get('reason', '')}")
            console.print("\n")
            
        # 在处理函数中显示提示信息
        console.print("\n[bold yellow]是否确认最终文本？[/bold yellow]")
        
        confirm = Prompt.ask("确认", choices=['yes', 'no'], default='yes')
        
        if confirm == 'yes':
            # 用户确认最终文本，结束流程
            return {
                'confirmed': True,
                'is_complete': True,
                'current_step': "end",
                'text': text,  # 保留确认的文本
                'final_revision': {  # 添加最终修订记录
                    'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                    'reason': "用户最终确认"
                }
            }
        else:
            # 用户不确认，返回到润色阶段
            return {
                'confirmed': False,
                'current_step': "human_review",
                'text': text  # 保留当前文本
            }
    
    register_interrupt_handler('final_confirm', final_confirm_handler)

# 主函数
def main():
    console.print(Panel.fit(
        "[bold cyan]Human-in-the-Loop 文本处理工作流[/bold cyan]\n"
        "这是一个交互式的文本处理系统，结合了AI分析和人工审核。",
        title="🤖 HITL Agent"
    ))
    
    # 注册中断处理器
    register_review_interrupt_handler()
    register_final_confirm_interrupt_handler()
    asyncio.run(run_hitl_workflow())

if __name__ == "__main__":
    main()
